Skip to content

Iterative CAGRA-Q - #1810

Open
irina-resh-nvda wants to merge 75 commits into
NVIDIA:release/26.10from
irina-resh-nvda:iterative_cagra_q
Open

Iterative CAGRA-Q#1810
irina-resh-nvda wants to merge 75 commits into
NVIDIA:release/26.10from
irina-resh-nvda:iterative_cagra_q

Conversation

@irina-resh-nvda

@irina-resh-nvda irina-resh-nvda commented Feb 16, 2026

Copy link
Copy Markdown
Contributor

Build CAGRA on PQ datasets with Iterative CAGRA-Q

Iterative cagra graph construction using CAGRA-Q search.

This PR improves the iterative CAGRA build method by enabling PQ compression: the dataset is compressed before the iterative search starts, and CAGRA-Q is used to iteratively update the KNN graph.

This is the first time we are introducing building CAGRA on (PQ) quantized datasets directly.

This PR also adds support for C and Python APIs for creating a PQ dataset and building a cagra graph on it. We also add support to update_dataset from dense->PQ or PQ->dense. Serializing the index is also supported.

Example Runnable workflow

import cupy as cp

from cuvs.common import make_device_pq_dataset
from cuvs.neighbors import cagra, hnsw
from cuvs.preprocessing.quantize import pq

dataset = cp.random.random((5000, 64), dtype=cp.float32)
pq_dataset = make_device_pq_dataset(
    pq.PQDatasetParams(pq_dim=32), dataset
)

index = cagra.build(
    cagra.IndexParams(build_algo="iterative_cagra_search"), pq_dataset
)
cagra.save("index.bin", index, include_dataset=False)

loaded_index = cagra.Index()
cagra.load(loaded_index, "index.bin")
cagra.update_dataset(loaded_index, dataset)

queries = dataset[:100]
cagra.search(cagra.SearchParams(), loaded_index, queries, 10)

hnsw_index = hnsw.from_cagra(hnsw.IndexParams(), loaded_index)
hnsw.search(hnsw.SearchParams(), hnsw_index, cp.asnumpy(queries), 10)

@copy-pr-bot

copy-pr-bot Bot commented Feb 16, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

… search

- Configurable growth-phase in-build search params (itopk_size, search_width,
  max_iterations) and internal/smem dtype; itopk auto-forced on the final
  full-size iteration.
- Decouple compression params used during iterative construction from the
  target index compression.
- Add shuffle_dataset option; fix out-of-bounds access from the in-place raft
  gather by switching to an out-of-place gather.
…around)

The shuffle_dataset path used an out-of-place gather into a temporary buffer to
work around an illegal memory access in raft's in-place gather overload when
n_rows * row_len exceeded 2^31 (32-bit index overflow).

That bug is now fixed upstream in raft (NVIDIA/raft#3059, closes #3055), which
the cuvs raft pin now includes. Revert to the in-place gather to drop the extra
full-size temporary allocation and copy.
@irina-resh-nvda
irina-resh-nvda marked this pull request as ready for review July 15, 2026 08:48
@irina-resh-nvda
irina-resh-nvda requested review from a team as code owners July 15, 2026 08:48
@aamijar aamijar changed the title Iterative cagra q Iterative graph build using CAGRA-Q search Jul 22, 2026
@aamijar aamijar changed the title Iterative graph build using CAGRA-Q search Iterative CAGRA-Q Jul 22, 2026
@tarang-jain
tarang-jain requested a review from a team as a code owner September 9, 2026 22:16
@tarang-jain

Copy link
Copy Markdown
Contributor

/ok to test 248e028

Comment thread cpp/src/neighbors/detail/dataset_serialize.hpp Outdated
Comment thread c/src/neighbors/cagra.cpp Outdated
Comment thread c/src/neighbors/cagra.cpp Outdated
Comment thread cpp/include/cuvs/neighbors/cagra.hpp Outdated
Comment thread cpp/include/cuvs/neighbors/cagra.hpp Outdated
Comment thread cpp/include/cuvs/neighbors/cagra.hpp Outdated
Comment thread cpp/src/neighbors/detail/cagra/cagra_build.cuh
Comment thread cpp/cmake/patches/faiss_override.json Outdated
Comment thread c/tests/CMakeLists.txt Outdated
@aamijar

aamijar commented Sep 10, 2026

Copy link
Copy Markdown
Member

/ok to test 100db4f

@irina-resh-nvda

Copy link
Copy Markdown
Contributor Author

It looks good to me

@aamijar

aamijar commented Sep 11, 2026

Copy link
Copy Markdown
Member

/ok to test 19f7397

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

feature request New feature or request non-breaking Introduces a non-breaking change

Projects

Status: In Progress

Development

Successfully merging this pull request may close these issues.

9 participants